"""LLM Plays Pokemon LLM客户端模块"""

import os
import logging
from typing import Dict, Any, Optional, List

logger = logging.getLogger(__name__)


class LLMClient:
    """LLM客户端类，负责与大型语言模型API交互"""
    
    def __init__(self, provider: str, model: str, api_key_env: str, max_tokens: int, base_url: str = "", tools: Optional[List[Dict[str, Any]]] = None):
        """初始化LLM客户端
        
        Args:
            provider: LLM提供商，如'openai'
            model: 模型名称
            api_key_env: API密钥值或环境变量名
            max_tokens: 最大生成令牌数
            base_url: API基础URL，为空则使用默认URL
            tools: 可用的函数工具列表，用于函数调用
        """
        self.provider = provider
        self.model = model
        # 直接使用api_key_env作为API密钥，不再从环境变量获取
        self.api_key = api_key_env
        self.max_tokens = max_tokens
        self.base_url = base_url
        self.tools = tools or []
        
        logger.info(f"LLM客户端初始化完成，提供商: {provider}, 模型: {model}")
    
    def generate(self, prompt: str) -> Any:
        """生成文本响应或函数调用
        
        Args:
            prompt: 提示词
            
        Returns:
            模型生成的响应文本或包含函数调用的消息对象
        """
        logger.info(f"向{self.provider}发送请求，使用模型: {self.model}")
        
        try:
            # 只支持OpenAI API调用
            if self.provider.lower() == "openai":
                return self._call_openai_api(prompt)
            else:
                logger.error(f"不支持的LLM提供商: {self.provider}")
                return "错误：不支持的LLM提供商"
        except Exception as e:
            logger.exception(f"调用LLM API失败: {e}")
            return f"错误：调用LLM API失败 - {str(e)}"
    
    def _call_openai_api(self, prompt: str) -> str:
        """调用OpenAI API
        
        Args:
            prompt: 提示词
            
        Returns:
            模型生成的响应文本
        """
        try:
            # 这里需要实现具体的OpenAI API调用
            # 简化版本，实际实现需要使用openai库
            import openai
            
            # 创建客户端，如果有自定义base_url则使用
            client_kwargs = {"api_key": self.api_key}
            if self.base_url:
                client_kwargs["base_url"] = self.base_url
                
            client = openai.OpenAI(**client_kwargs)
            
            # 解析提示词中的系统提示、用户消息和助手消息
            messages = self._parse_prompt_to_messages(prompt)
            
            # 准备API调用参数
            api_params = {
                "model": self.model,
                "messages": messages,
                "max_tokens": self.max_tokens
            }
            
            # 如果有工具定义，添加到API调用参数中
            if self.tools:
                api_params["tools"] = self.tools
                api_params["tool_choice"] = "auto"
            
            response = client.chat.completions.create(**api_params)
            
            # 处理响应
            message = response.choices[0].message
            
            # 如果有工具调用，返回完整消息对象而不仅仅是内容
            if hasattr(message, 'tool_calls') and message.tool_calls:
                logger.info(f"LLM返回了工具调用: {len(message.tool_calls)}个")
                return message
            
            # 否则返回消息内容
            return message.content
        except ImportError:
            logger.error("未安装openai库，请使用'pip install openai'安装")
            return "错误：未安装openai库"
        except Exception as e:
            logger.exception(f"调用OpenAI API失败: {e}")
            return f"错误：调用OpenAI API失败 - {str(e)}"
    
    def _parse_prompt_to_messages(self, prompt: str) -> list:
        """将提示词解析为消息列表
        
        Args:
            prompt: 提示词
            
        Returns:
            消息列表
        """
        # 简单的解析逻辑，实际实现可能需要更复杂的解析
        lines = prompt.split("\n")
        messages = []
        
        # 提取系统提示
        system_content = ""
        user_content = ""
        assistant_content = ""
        current_role = "system"
        
        for line in lines:
            if line.startswith("用户:"):
                if system_content and current_role == "system":
                    messages.append({"role": "system", "content": system_content.strip()})
                    system_content = ""
                if user_content and current_role == "user":
                    messages.append({"role": "user", "content": user_content.strip()})
                    user_content = ""
                if assistant_content and current_role == "assistant":
                    messages.append({"role": "assistant", "content": assistant_content.strip()})
                    assistant_content = ""
                current_role = "user"
                user_content += line[3:].strip() + "\n"
            elif line.startswith("助手:"):
                if system_content and current_role == "system":
                    messages.append({"role": "system", "content": system_content.strip()})
                    system_content = ""
                if user_content and current_role == "user":
                    messages.append({"role": "user", "content": user_content.strip()})
                    user_content = ""
                if assistant_content and current_role == "assistant":
                    messages.append({"role": "assistant", "content": assistant_content.strip()})
                    assistant_content = ""
                current_role = "assistant"
                assistant_content += line[3:].strip() + "\n"
            elif line.startswith("系统:"):
                if system_content and current_role == "system":
                    messages.append({"role": "system", "content": system_content.strip()})
                    system_content = ""
                if user_content and current_role == "user":
                    messages.append({"role": "user", "content": user_content.strip()})
                    user_content = ""
                if assistant_content and current_role == "assistant":
                    messages.append({"role": "assistant", "content": assistant_content.strip()})
                    assistant_content = ""
                current_role = "system"
                system_content += line[3:].strip() + "\n"
            else:
                if current_role == "system":
                    system_content += line + "\n"
                elif current_role == "user":
                    user_content += line + "\n"
                elif current_role == "assistant":
                    assistant_content += line + "\n"
        
        # 添加最后一个角色的内容
        if system_content and current_role == "system":
            messages.append({"role": "system", "content": system_content.strip()})
        if user_content and current_role == "user":
            messages.append({"role": "user", "content": user_content.strip()})
        if assistant_content and current_role == "assistant":
            messages.append({"role": "assistant", "content": assistant_content.strip()})
        
        # 如果没有消息，添加一个默认的用户消息
        if not messages:
            messages.append({"role": "user", "content": prompt})
        
        return messages